# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/5 16:29 
@Author  : ZhangShenao 
@File    : multi_query_rrf.py 
@Desc    : 
"""
from typing import List

from langchain.retrievers import MultiQueryRetriever
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.load import dumps, loads


class ReciprocalRankFusionRetriever(MultiQueryRetriever):
    """Enhancement of MultiQueryRetriever

    Retrieve docs for each query. Using the RRF algorithm,
    sort and merge all retrieved documents.

    Reference paper: https://plg.uwaterloo.ca/~gvcormac/cormacksigir09-rrf.pdf
    """

    c: int = 60
    """A constant added to the rank, controlling the balance between the importance
    of high-ranked items and the consideration given to lower-ranked items.
    Default is 60.
    """

    top_k: int = 3
    """Return final top-k documents ranked by rrf score"""

    def __init__(self, top_k: int = 3, c: int = 60, **kwargs):
        super().__init__(**kwargs)
        self.top_k = top_k
        self.c = c

    def retrieve_documents(
            self, queries: List[str], run_manager: CallbackManagerForRetrieverRun
    ) -> List[List[Document]]:
        """Run all LLM generated queries.

        Args:
            queries: query list

        Returns:
            Nested list of retrieved Documents
            Group by user query
        """

        documents = []
        for query in queries:
            docs = self.retriever.invoke(
                query, config={"callbacks": run_manager.get_child()}
            )
            documents.append(docs)
        return documents

    def unique_union(self, documents: List[List[Document]]) -> List[Document]:
        """Get unique Documents.

        Args:
            documents: List of retrieved Documents

        Returns:
            List of unique retrieved Documents
        """

        return self._rank_fusion(documents)

    def _rank_fusion(self, documents: List[List[Document]]) -> List[Document]:
        """Sort and merge all retrieved documents by the RRF algorithm"""

        fused_scores = {}

        for docs in documents:
            for rank, doc in enumerate(docs):
                doc_str = dumps(doc)
                if doc_str not in fused_scores:
                    fused_scores[doc_str] = 0
                fused_scores[doc_str] += 1 / (rank + self.c)

        reranked_results = [
            (loads(doc), score)
            for doc, score in sorted(fused_scores.items(), key=lambda x: x[1], reverse=True)
        ]

        return [item[0] for item in reranked_results[:self.top_k]]
